//! 并发工具

use crossbeam::channel;
use rayon::prelude::*;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

/// 线程池执行器
#[derive(Debug)]
pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<channel::Sender<Job>>,
}

type Job = Box<dyn FnOnce() + Send + 'static>;

impl ThreadPool {
    /// 创建指定大小的线程池
    /// 
    /// # 示例
    /// 
    /// ```
    /// use rust_utils::concurrent_utils::ThreadPool;
    /// 
    /// let pool = ThreadPool::new(4);
    /// assert!(pool.is_ok());
    /// ```
    pub fn new(size: usize) -> Result<ThreadPool, &'static str> {
        if size == 0 {
            return Err("线程池大小必须大于0");
        }

        let (sender, receiver) = channel::unbounded();
        let receiver = Arc::new(Mutex::new(receiver));
        let mut workers = Vec::with_capacity(size);

        for id in 0..size {
            workers.push(Worker::new(id, Arc::clone(&receiver)));
        }

        Ok(ThreadPool {
            workers,
            sender: Some(sender),
        })
    }

    /// 在线程池中执行任务
    /// 
    /// # 示例
    /// 
    /// ```
    /// use rust_utils::concurrent_utils::ThreadPool;
    /// use std::sync::{Arc, Mutex};
    /// use std::time::Duration;
    /// use std::thread;
    /// 
    /// let pool = ThreadPool::new(4).unwrap();
    /// let counter = Arc::new(Mutex::new(0));
    /// let counter_clone = Arc::clone(&counter);
    /// 
    /// pool.execute(move || {
    ///     let mut num = counter_clone.lock().unwrap();
    ///     *num += 1;
    /// });
    /// 
    /// // 等待任务完成
    /// thread::sleep(Duration::from_millis(100));
    /// assert_eq!(*counter.lock().unwrap(), 1);
    /// ```
    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + Send + 'static,
    {
        let job = Box::new(f);
        self.sender.as_ref().unwrap().send(job).unwrap();
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        drop(self.sender.take());

        for worker in &mut self.workers {
            println!("关闭工作线程 {}", worker.id);

            if let Some(thread) = worker.thread.take() {
                thread.join().unwrap();
            }
        }
        println!("所有工作线程已关闭");
    }
}

#[derive(Debug)]
struct Worker {
    id: usize,
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    fn new(id: usize, receiver: Arc<Mutex<channel::Receiver<Job>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver.lock().unwrap().recv();

            match message {
                Ok(job) => {
                    println!("工作线程 {} 执行任务", id);
                    job();
                }
                Err(_) => {
                    println!("工作线程 {} 断开连接", id);
                    break;
                }
            }
        });

        Worker {
            id,
            thread: Some(thread),
        }
    }
}

/// 并行计算Vec中所有元素的和
/// 
/// # 示例
/// 
/// ```
/// use rust_utils::concurrent_utils::parallel_sum;
/// 
/// let data = vec![1, 2, 3, 4, 5];
/// let sum = parallel_sum(&data);
/// assert_eq!(sum, 15);
/// ```
pub fn parallel_sum(data: &[i32]) -> i32 {
    data.par_iter().sum()
}

/// 并行映射函数
/// 
/// # 示例
/// 
/// ```
/// use rust_utils::concurrent_utils::parallel_map;
/// 
/// let data = vec![1, 2, 3, 4, 5];
/// let result: Vec<i32> = parallel_map(&data, |x| x * 2);
/// assert_eq!(result, vec![2, 4, 6, 8, 10]);
/// ```
pub fn parallel_map<T, F, R>(data: &[T], f: F) -> Vec<R>
where
    T: Sync + Clone,
    R: Send,
    F: Fn(T) -> R + Sync,
{
    data.par_iter().map(|x| f(x.clone())).collect()
}

/// 并行过滤函数
/// 
/// # 示例
/// 
/// ```
/// use rust_utils::concurrent_utils::parallel_filter;
/// 
/// let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
/// let result: Vec<i32> = parallel_filter(&data, |&x| x % 2 == 0);
/// assert_eq!(result, vec![2, 4, 6, 8, 10]);
/// ```
pub fn parallel_filter<T, F>(data: &[T], predicate: F) -> Vec<T>
where
    T: Sync + Clone + Send,
    F: Fn(&T) -> bool + Sync,
{
    data.par_iter().filter(|&x| predicate(x)).cloned().collect()
}

/// 带超时的通道发送
/// 
/// # 示例
/// 
/// ```
/// use rust_utils::concurrent_utils::send_with_timeout;
/// use crossbeam::channel;
/// use std::time::Duration;
/// 
/// let (s, r) = channel::unbounded();
/// let result = send_with_timeout(s, 42, Duration::from_millis(100));
/// assert!(result.is_ok());
/// assert_eq!(r.recv().unwrap(), 42);
/// ```
pub fn send_with_timeout<T: Send>(
    sender: channel::Sender<T>,
    msg: T,
    timeout: Duration,
) -> Result<(), channel::SendTimeoutError<T>> {
    sender.send_timeout(msg, timeout)
}

/// 带超时的通道接收
/// 
/// # 示例
/// 
/// ```
/// use rust_utils::concurrent_utils::recv_with_timeout;
/// use crossbeam::channel;
/// use std::time::Duration;
/// use std::thread;
/// 
/// let (s, r) = channel::unbounded();
/// 
/// thread::spawn(move || {
///     thread::sleep(Duration::from_millis(50));
///     s.send(42).unwrap();
/// });
/// 
/// let result = recv_with_timeout(r, Duration::from_millis(100));
/// assert!(result.is_ok());
/// assert_eq!(result.unwrap(), 42);
/// ```
pub fn recv_with_timeout<T>(
    receiver: channel::Receiver<T>,
    timeout: Duration,
) -> Result<T, channel::RecvTimeoutError> {
    receiver.recv_timeout(timeout)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};
    use std::thread;
    use std::time::Duration;

    #[test]
    fn test_thread_pool() {
        let pool = ThreadPool::new(4);
        assert!(pool.is_ok());

        let counter = Arc::new(Mutex::new(0));
        let counter_clone = Arc::clone(&counter);

        pool.unwrap().execute(move || {
            let mut num = counter_clone.lock().unwrap();
            *num += 1;
        });

        // 等待任务完成
        thread::sleep(Duration::from_millis(100));
        assert_eq!(*counter.lock().unwrap(), 1);
    }

    #[test]
    fn test_parallel_sum() {
        let data = vec![1, 2, 3, 4, 5];
        let sum = parallel_sum(&data);
        assert_eq!(sum, 15);
    }

    #[test]
    fn test_parallel_map() {
        let data = vec![1, 2, 3, 4, 5];
        let result: Vec<i32> = parallel_map(&data, |x| x * 2);
        assert_eq!(result, vec![2, 4, 6, 8, 10]);
    }

    #[test]
    fn test_parallel_filter() {
        let data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        let result: Vec<i32> = parallel_filter(&data, |&x| x % 2 == 0);
        assert_eq!(result, vec![2, 4, 6, 8, 10]);
    }

    #[test]
    fn test_send_with_timeout() {
        let (s, r) = channel::unbounded();
        let result = send_with_timeout(s, 42, Duration::from_millis(100));
        assert!(result.is_ok());
        assert_eq!(r.recv().unwrap(), 42);
    }

    #[test]
    fn test_recv_with_timeout() {
        let (s, r) = channel::unbounded();

        thread::spawn(move || {
            thread::sleep(Duration::from_millis(50));
            s.send(42).unwrap();
        });

        let result = recv_with_timeout(r, Duration::from_millis(100));
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 42);
    }
}